def parse_file(file_path):
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    current_id = 1
    for line in lines:
        line = line.strip()
        if not line or line.startswith('`') or 'COMMENT' not in line:
            continue

        # 提取字段名和内容
        parts = line.split('COMMENT')
        if len(parts) == 2:
            content = parts[1].strip().strip("'")
            animal_name = f"Animal {current_id}"
            
            # 检查内容是否为空
            if content:
                embedding = model.encode(content).tolist()
                data.append([current_id, animal_name, content, embedding])
                current_id += 1

    return data

# 连接 Milvus
MILVUS_HOST = '124.220.19.41'
MILVUS_PORT = '19530'

from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
from sentence_transformers import SentenceTransformer

connections.connect(alias="default", host=MILVUS_HOST, port=MILVUS_PORT)

# 创建 collection，维度修改为 384
collection_name = "animal_plant_identification"

if utility.has_collection(collection_name):
    utility.drop_collection(collection_name)

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
    FieldSchema(name="animal_name", dtype=DataType.VARCHAR, max_length=100),
    FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=4096),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=384)   # 修复为模型维度
]

schema = CollectionSchema(fields, description="Animal and plant identification data")
collection = Collection(name=collection_name, schema=schema)

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri="http://localhost:19530",
    token="root:Milvus"
)

# 3.1. Create schema
schema = MilvusClient.create_schema(
    auto_id=False,
    enable_dynamic_field=True,
)

# 加载 384维模型
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

# 数据解析
def parse_file(file_path):
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    current_id = 1
    for line in lines:
        line = line.strip()
        if not line or line.startswith('`') or 'COMMENT' not in line:
            continue

        parts = line.split('COMMENT')
        if len(parts) == 2:
            content = parts[1].strip().strip("'")
            animal_name = f"Animal {current_id}"
            
            # 生成 384维向量
            if content:
                embedding = model.encode(content).tolist()
                data.append([current_id, animal_name, content, embedding])
                current_id += 1

    return data

# 插入数据
data = parse_file("c:/Users/大糊/Desktop/animal-plant-identification.txt")

if len(data) == 0:
    print("❌ 数据解析失败！")
else:
    insert_data = list(zip(*data))
    collection.insert(insert_data)
    collection.flush()
    collection.load()
    print("✅ 数据成功插入 Milvus！")
